from __future__ import annotations

import os, json
import logging

from typing import List, Optional, Any, Dict

from fastapi import APIRouter, HTTPException

from pydantic import BaseModel, Field

from openai import OpenAI

from ..config import get_settings
from ..services.mcp_session import mcp


logger = logging.getLogger(__name__)



router = APIRouter(prefix="/ai", tags=["ai"])



class ChatReq(BaseModel):
    """
    AI 对话请求模型
    
    支持通过自然语言与 AI 助手交互，访问 Hummingbot MCP 工具来管理交易账户、执行交易和分析市场数据。
    参考文档: https://hummingbot.org/mcp/tools/
    """
    prompt: str = Field(
        ...,
        description="用户自然语言指令。例如：'查看我的账户余额'、'BTC 当前价格是多少？'、'帮我买入 0.1 BTC'",
        examples=["查看我的账户余额", "BTC 当前价格是多少？", "帮我买入 0.1 BTC"]
    )
    
    execute: bool = Field(
        False,
        description="是否允许执行可能的下单类动作（全局安全开关）。\n"
                  "- `false`（默认）: 拦截所有交易类操作（下单、取消订单、平仓等），仅返回模拟结果\n"
                  "- `true`: 允许执行实际的交易操作\n"
                  "⚠️ 注意：设置为 true 前请确保已理解操作风险",
        examples=[False, True]
    )
    
    allowed_tool_categories: Optional[List[str]] = Field(
        None,
        description="允许使用的工具类别白名单。为空时允许所有类别。\n\n"
                   "可选值：\n"
                   "- `account`: 账户工具（get_accounts, get_account_balance）\n"
                   "- `portfolio`: 投资组合工具（get_portfolio_balances, get_portfolio_performance）\n"
                   "- `trading`: 交易工具（place_order, cancel_order, get_open_orders, get_order_history）\n"
                   "- `position`: 持仓工具（get_positions, close_position, get_position_history）\n"
                   "- `market_data`: 市场数据工具（get_ticker, get_orderbook, get_funding_rates）\n\n"
                   "示例：仅允许查询账户和市场数据，禁止交易操作，可设置为 `['account', 'market_data']`",
        examples=[
            None,
            ["account", "market_data"],
            ["account", "portfolio", "market_data"]
        ]
    )
    
    blocked_tools: Optional[List[str]] = Field(
        None,
        description="被阻止的工具列表（工具名称黑名单）。优先级高于 allowed_tool_categories。\n\n"
                   "常用工具名称：\n"
                   "- `place_order`: 下单\n"
                   "- `cancel_order`: 取消订单\n"
                   "- `close_position`: 平仓\n"
                   "- `get_positions`: 查询持仓\n\n"
                   "示例：`['place_order', 'cancel_order']` 将阻止下单和取消订单操作",
        examples=[
            None,
            ["place_order"],
            ["place_order", "cancel_order", "close_position"]
        ]
    )
    
    max_tool_rounds: int = Field(
        2,
        ge=1,
        le=10,
        description="最大工具调用轮次。\n"
                   "- 每轮可能包含多个工具调用\n"
                   "- 设置为较小值（如 1-2）可限制操作步骤，提高安全性\n"
                   "- 复杂查询可能需要更多轮次（如 3-5）\n"
                   "- 范围：1-10",
        examples=[1, 2, 5]
    )
    
    temperature: Optional[float] = Field(
        None,
        ge=0.0,
        le=2.0,
        description="模型温度参数，控制响应的随机性和创造性。\n"
                   "- 较低值（0.0-0.3）：更确定性、更保守的响应，适合交易操作\n"
                   "- 中等值（0.5-0.7）：平衡的响应，推荐用于一般查询\n"
                   "- 较高值（0.8-2.0）：更创造性、更多样化的响应\n"
                   "- 为空时使用模型默认值（推荐）\n"
                   "- 范围：0.0-2.0",
        examples=[None, 0.3, 0.7, 1.0]
    )
    
    max_tokens: Optional[int] = Field(
        None,
        ge=1,
        description="最大生成 token 数限制。\n"
                   "- 控制 AI 响应的最大长度\n"
                   "- 为空时不限制（推荐）\n"
                   "- 建议值：500-2000（短响应），2000-4000（详细分析）\n"
                   "- 最小：1",
        examples=[None, 500, 1000, 2000]
    )
    
    system_prompt: Optional[str] = Field(
        None,
        description="自定义系统提示词，用于指导 AI 的行为和响应风格。\n\n"
                   "为空时使用默认提示词（推荐），默认提示词包括：\n"
                   "- 工具类别说明\n"
                   "- 操作安全指南\n"
                   "- 响应格式要求\n\n"
                   "示例：可自定义提示词来强调特定的交易策略或风险偏好",
        examples=[None, "你是一个保守的交易助手，执行任何操作前都需要用户明确确认。"]
    )
    
    stream: bool = Field(
        False,
        description="是否启用流式响应（Server-Sent Events）。\n"
                  "- `false`（默认）：返回完整响应\n"
                  "- `true`: 流式返回响应（暂不支持，保留用于未来扩展）\n"
                  "⚠️ 当前版本暂不支持流式响应，请保持为 false",
        examples=[False]
    )



class ChatResp(BaseModel):
    """
    AI 对话响应模型
    
    包含 AI 的文本回复、调用的工具列表、Token 使用统计等信息。
    """
    text: str = Field(
        ...,
        description="AI 的文本回复内容。包含对用户问题的回答和操作结果的解释。",
        examples=["根据查询结果，您当前在 Binance 交易所的 BTC 余额为 1.5 BTC。"]
    )
    
    tool_invocations: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="本次对话中调用的工具列表。每个工具调用包含：\n"
                  "- `name`: 工具名称（如 'get_account_balance'）\n"
                  "- `args`: 工具参数（JSON 对象）\n"
                  "- `result`: 工具执行结果（字符串或对象）\n\n"
                  "示例：`[{'name': 'get_account_balance', 'args': {'account_name': 'binance'}, 'result': '...'}]`\n\n"
                  "⚠️ 重要：参数名称必须与工具定义完全一致。例如：\n"
                  "- `get_account_balance` 使用 `account_name`（不是 `account`）\n"
                  "- `get_ticker` 使用 `exchange` 和 `symbol`（不是 `trading_pair`）",
        examples=[
            [],
            [
                {
                    "name": "get_account_balance",
                    "args": {"account_name": "binance"},
                    "result": "BTC: 1.5, USDT: 10000"
                }
            ]
        ]
    )
    
    tokens_used: Optional[Dict[str, Any]] = Field(
        None,
        description="Token 使用统计信息（如果可用）。包含：\n"
                   "- `prompt_tokens`: 输入提示词消耗的 token 数\n"
                   "- `completion_tokens`: 生成响应消耗的 token 数\n"
                   "- `total_tokens`: 总 token 数（prompt_tokens + completion_tokens）\n\n"
                   "如果 API 响应中不包含使用统计，此字段为 null。",
        examples=[
            None,
            {
                "prompt_tokens": 150,
                "completion_tokens": 200,
                "total_tokens": 350
            }
        ]
    )
    
    finish_reason: Optional[str] = Field(
        None,
        description="对话完成的原因。可能的值：\n"
                "- `stop`: 正常完成，AI 生成完整的回复\n"
                "- `length`: 达到最大 token 限制而截断\n"
                "- `tool_calls`: 需要调用工具来继续（通常不会出现在最终响应中）\n"
                "- `null`: 未知或其他原因\n\n"
                "如果 API 响应中不包含此信息，此字段为 null。",
        examples=[None, "stop", "length"]
    )



def _get_tool_category(tool_name: str) -> str:
    """根据工具名称推断其类别"""
    name_lower = tool_name.lower()
    
    if "account" in name_lower:
        return "account"
    elif "portfolio" in name_lower:
        return "portfolio"
    elif any(x in name_lower for x in ["order", "trade", "buy", "sell"]):
        return "trading"
    elif "position" in name_lower:
        return "position"
    elif any(x in name_lower for x in [
        "ticker",
        "price",
        "orderbook",
        "candle",
        "kline",
        "funding",
        "market",
    ]):
        return "market_data"
    else:
        return "other"


def _filter_tools(
    tools: List[Any],
    allowed_categories: Optional[List[str]] = None,
    blocked_tools: Optional[List[str]] = None
) -> List[Any]:
    """根据权限设置过滤工具"""
    # 验证并清理 allowed_categories，只保留有效类别
    valid_categories = {"account", "portfolio", "trading", "position", "market_data", "other"}
    if allowed_categories:
        allowed_categories = [cat for cat in allowed_categories if cat in valid_categories]
        if not allowed_categories:
            logger.warning("allowed_tool_categories 中没有有效类别，将允许所有工具")
            allowed_categories = None
    
    if not allowed_categories and not blocked_tools:
        return tools
    
    filtered = []
    for tool in tools:
        tool_name = tool.name
        
        # 检查是否在阻止列表中
        if blocked_tools and tool_name in blocked_tools:
            logger.debug(f"工具 {tool_name} 在阻止列表中，已过滤")
            continue
        
        # 检查类别是否允许
        if allowed_categories:
            category = _get_tool_category(tool_name)
            if category not in allowed_categories:
                logger.debug(f"工具 {tool_name} 类别 {category} 不在允许列表中，已过滤")
                continue
        
        filtered.append(tool)
        logger.debug(f"工具 {tool_name} 已通过过滤")
    
    return filtered


def _mcp_tools_to_openai(tools):

    """把 MCP 工具转换成 OpenAI function-tools 规范"""

    out = []

    for t in tools:

        out.append({

            "type": "function",

            "function": {

                "name": t.name,

                "description": t.description or "",

                "parameters": t.inputSchema or {"type": "object", "properties": {}}

            }

        })

    return out



def _flatten_mcp_result(res) -> str:

    # MCP 返回的内容是一个 content 列表，这里简单抽取 text

    parts = []

    for c in getattr(res, "content", []) or []:

        if getattr(c, "type", "") == "text":

            parts.append(getattr(c, "text", ""))

    return "\n".join([p for p in parts if p])



@router.post(
    "/chat",
    response_model=ChatResp,
    summary="AI 对话接口",
    description="""
与 AI 助手进行对话，支持通过自然语言访问 Hummingbot MCP 工具。

## 功能特性

- **自然语言交互**: 使用自然语言描述需求，AI 自动选择合适的工具执行
- **工具类别支持**: 
  - Account Tools: 账户和余额查询
  - Portfolio Tools: 投资组合分析
  - Trading Tools: 订单管理（下单、取消等）
  - Position Tools: 持仓管理
  - Market Data Tools: 市场数据查询

## 安全控制

- **执行开关**: `execute=false` 时拦截所有交易操作，仅返回模拟结果
- **工具过滤**: 通过 `allowed_tool_categories` 和 `blocked_tools` 精确控制可用工具
- **轮次限制**: `max_tool_rounds` 限制操作步骤数量

## 使用示例

**查询账户余额**:
```json
{
  "prompt": "查看我在 Binance 的账户余额",
  "execute": false,
  "allowed_tool_categories": ["account", "market_data"]
}
```

**查询市场数据**:
```json
{
  "prompt": "BTC 当前价格是多少？",
  "execute": false
}
```

**模拟下单（不实际执行）**:
```json
{
  "prompt": "帮我买入 0.1 BTC",
  "execute": false
}
```

参考文档: https://hummingbot.org/mcp/tools/
    """,
    responses={
        200: {
            "description": "成功返回 AI 回复和工具调用结果",
            "content": {
                "application/json": {
                    "example": {
                        "text": "根据查询，您当前在 Binance 的 BTC 余额为 1.5 BTC。",
                        "tool_invocations": [
                            {
                                "name": "get_account_balance",
                                "args": {"account_name": "binance"},
                                "result": "BTC: 1.5, USDT: 10000"
                            }
                        ],
                        "tokens_used": {
                            "prompt_tokens": 150,
                            "completion_tokens": 200,
                            "total_tokens": 350
                        },
                        "finish_reason": "stop"
                    }
                }
            }
        },
        500: {
            "description": "服务器错误，可能是 MCP 服务未配置或 DeepSeek API 密钥未设置"
        }
    }
)
async def chat(req: ChatReq) -> ChatResp:

    # 1) 确保 MCP 会话可用

    await mcp.start()
    
    if not mcp.enabled:
        return ChatResp(
            text="MCP 服务未配置或启动失败。请设置 MCP_COMMAND 和 MCP_ARGS_JSON 环境变量。",
            tool_invocations=[]
        )
    
    try:
        tools = await mcp.list_tools()
        
        # 应用工具过滤
        tools = _filter_tools(
            tools,
            allowed_categories=req.allowed_tool_categories,
            blocked_tools=req.blocked_tools
        )
        
        if not tools:
            return ChatResp(
                text="没有可用的工具。请检查 allowed_tool_categories 和 blocked_tools 参数。",
                tool_invocations=[]
            )
        
        openai_tools = _mcp_tools_to_openai(tools)
        
        # 记录工具信息用于调试
        logger.info(f"过滤后的工具数量: {len(tools)}, 转换为 OpenAI 格式: {len(openai_tools)}")
        if tools:
            tool_names = [t.name for t in tools]
            logger.info(f"可用工具列表: {tool_names}")
        
    except RuntimeError as e:
        return ChatResp(
            text=f"MCP 服务不可用: {str(e)}",
            tool_invocations=[]
        )



    # 2) DeepSeek（OpenAI 兼容）客户端
    settings = get_settings()
    
    if not settings.deepseek_api_key:
        raise HTTPException(
            status_code=500,
            detail="DEEPSEEK_API_KEY 未配置。请在 .env 文件中设置 DEEPSEEK_API_KEY。"
        )

    client = OpenAI(
        api_key=settings.deepseek_api_key,
        base_url=settings.deepseek_base_url,
    )
    model = settings.deepseek_model



    # 构建消息列表
    messages: List[Dict[str, Any]] = []
    
    # 添加系统提示词（如果提供）
    if req.system_prompt:
        messages.append({"role": "system", "content": req.system_prompt})
    else:
        # 默认系统提示词，基于 Hummingbot MCP Tools 文档
        tool_count = len(openai_tools) if openai_tools else 0
        tool_names = [t["function"]["name"] for t in openai_tools] if openai_tools else []
        
        default_system_prompt = f"""你是一个专业的交易助手，可以访问 Hummingbot MCP 工具来帮助用户管理交易账户、执行交易和分析市场数据。

当前可用工具数量: {tool_count}
可用工具列表: {', '.join(tool_names[:10])}{'...' if len(tool_names) > 10 else ''}

重要提示：
1. **必须使用工具来获取数据**：当用户询问账户余额、价格、持仓等信息时，你必须调用相应的工具来获取真实数据，而不是猜测或说"功能不可用"

工具使用方法（重要：注意参数名称）：
- **get_account_balance**: 查询特定账户余额
  * 参数: `account_name` (string, 必需) - 交易所账户名称，例如 "binance", "okx"
  * 示例: `{{"account_name": "binance"}}`
  
- **get_accounts**: 列出所有连接的交易所账户
  * 参数: 无
  
- **get_ticker**: 获取交易对当前价格
  * 参数: `exchange` (string, 必需) - 交易所名称, `symbol` (string, 必需) - 交易对符号，例如 "BTC-USDT"
  * 示例: `{{"exchange": "binance", "symbol": "BTC-USDT"}}`
  
- **get_positions**: 查询持仓
  * 参数: `exchange` (string, 可选) - 交易所名称，为空时查询所有交易所
  
- **get_open_orders**: 查询未完成订单
  * 参数: `exchange` (string, 可选) - 交易所名称
  
- **get_order_history**: 查询订单历史
  * 参数: `exchange` (string, 可选), `trading_pair` (string, 可选), `limit` (number, 可选)

工具类别说明：
- Account Tools (account): 管理交易所账户和余额（get_accounts, get_account_balance）
- Portfolio Tools (portfolio): 查看和分析投资组合（get_portfolio_balances, get_portfolio_performance）
- Trading Tools (trading): 执行买卖订单和管理订单（place_order, cancel_order, get_open_orders, get_order_history）
- Position Tools (position): 管理衍生品持仓（get_positions, close_position, get_position_history）
- Market Data Tools (market_data): 获取市场数据和价格信息（get_ticker, get_orderbook, get_funding_rates）

请始终：
1. **优先使用工具获取真实数据**，不要直接回答"功能不可用"
2. 在执行任何交易操作前，向用户确认意图和参数
3. 清晰地解释每个操作的含义和影响
4. 提供准确的数据和分析结果
5. 在出现错误时给出清晰的错误说明

根据用户的问题，必须选择合适的工具并执行相应的操作来获取真实数据。"""
        messages.append({"role": "system", "content": default_system_prompt})
    
    messages.append({"role": "user", "content": req.prompt})

    tool_invocations: List[Dict[str, Any]] = []



    # 3) 迭代执行（工具调用 -> 写回结果 -> 让模型给最终答案）

    for _ in range(max(req.max_tool_rounds, 1)):

        # 构建请求参数
        completion_params = {
            "model": model,
            "messages": messages,
            "tools": openai_tools,
            "tool_choice": "auto",
        }
        
        # 添加可选参数
        if req.temperature is not None:
            completion_params["temperature"] = req.temperature
        if req.max_tokens is not None:
            completion_params["max_tokens"] = req.max_tokens
        
        resp = client.chat.completions.create(**completion_params)

        choice = resp.choices[0]

        msg = choice.message



        if msg.tool_calls:
            tool_call_payloads = [tc.model_dump() for tc in msg.tool_calls]

            # 先把 assistant 带 tool_calls 的消息写回（一次）
            messages.append({
                "role": "assistant",
                "content": msg.content or "",
                "tool_calls": tool_call_payloads,
            })

            # 安全阀：如果 execute=False，直接把意图/参数解释出来，不真正执行下单等动作

            for call in msg.tool_calls:

                name = call.function.name

                args = json.loads(call.function.arguments or "{}")

                # 更全面的交易类操作检测
                trading_keywords = [
                    "place_order",
                    "cancel_order",
                    "close_position",
                    "buy",
                    "sell",
                ]
                name_lower = name.lower()
                is_trading_action = any(keyword in name_lower for keyword in trading_keywords)
                
                if not req.execute and is_trading_action:

                    # 不执行，返回解释

                    tool_invocations.append({"name": name, "args": args, "result": "[dry_run] 已拦截可能的交易操作"})

                    messages.append({

                        "role": "tool",

                        "tool_call_id": call.id,

                        "content": "[dry_run] 交易类操作已拦截，请确认 execute=true 后再试。"

                    })

                    continue



                # 真正执行 MCP 工具

                mcp_res = await mcp.call_tool(name, args)

                out_text = _flatten_mcp_result(mcp_res) or str([c.model_dump() for c in mcp_res.content])

                tool_invocations.append({"name": name, "args": args, "result": out_text})



                # 把工具结果按 OpenAI 规范写回
                messages.append({
                    "role": "tool",
                    "tool_call_id": call.id,
                    "content": out_text
                })
                # 如果还有后续工具调用需要处理，继续；否则结束当前循环
            continue



        # 没有工具调用，直接返回
        tokens_used = None
        if hasattr(resp, "usage"):
            tokens_used = {
                "prompt_tokens": resp.usage.prompt_tokens,
                "completion_tokens": resp.usage.completion_tokens,
                "total_tokens": resp.usage.total_tokens,
            } if resp.usage else None
        
        finish_reason = choice.finish_reason
        
        return ChatResp(
            text=msg.content or "",
            tool_invocations=tool_invocations,
            tokens_used=tokens_used,
            finish_reason=finish_reason
        )



    return ChatResp(
        text="未能得到模型有效回复，请稍后重试。",
        tool_invocations=tool_invocations,
        tokens_used=None,
        finish_reason=None
    )

